import matplotlib.pyplot as plt
import numpy as np
from collections import deque

# ========== IMC-PID 控制器 ==========
class IMCPID:
    def __init__(self, K=1.0, tau=50.0, theta=5.0, lam=40.0, dt=1.0):
        self.dt = dt
        self.lam = lam
        self.set_params(K, tau, theta, lam)

        # 控制器状态
        self.integral = 0.0
        self.prev_measurement = 0.0

        # 积分限幅
        self.int_min = -5.0
        self.int_max =  5.0

    def set_params(self, K, tau, theta, lam=None):
        if lam is not None:
            self.lam = lam

        eps = 1e-6
        K  = float(np.clip(K, 0.05, 50.0))
        tau = float(np.clip(tau, 1.0, 5_000.0))
        theta = float(np.clip(theta, 0.0, 1_000.0))

        self.K = K
        self.tau = tau
        self.theta = theta

        # IMC-PID 推荐公式
        self.Kp = (2.0 * tau + theta) / (K * (2.0 * self.lam + theta))
        self.Ti = tau + theta / 2.0
        self.Td = (tau * theta) / (2.0 * tau + theta + eps)

        # 限幅
        self.Kp = float(np.clip(self.Kp, 0.0, 10.0))
        self.Ti = float(np.clip(self.Ti, self.dt, 10_000.0))
        self.Td = float(np.clip(self.Td, 0.0, 1_000.0))

    def compute(self, setpoint, measurement):
        error = setpoint - measurement

        self.integral += error * self.dt
        self.integral = np.clip(self.integral, self.int_min, self.int_max)

        derivative = (self.prev_measurement - measurement) / self.dt

        output = ( self.Kp * error
                 + (self.Kp / self.Ti) * self.integral
                 - self.Kp * self.Td * derivative )

        self.prev_measurement = measurement
        return output


# ========== 在线RLS+延时估计 ==========
class OnlineRLSFOPDT:
    def __init__(self, dt=1.0, d_max=300, lam_corr=2000,
                 rls_lambda=0.995, init_var=1e3, retune_period=120,
                 smooth_alpha=0.2):
        self.dt = dt
        self.d_max = d_max
        self.corr_win = lam_corr
        self.rls_lambda = rls_lambda
        self.P = np.eye(2) * init_var
        self.theta = np.array([0.99, 0.001])
        self.d = 5
        self.retune_period = retune_period
        self.smooth_alpha = smooth_alpha

        self.y_hist = deque(maxlen=max(lam_corr, d_max+5))
        self.u_hist = deque(maxlen=max(lam_corr, d_max+5))

        self.estimates = []

        self.K_hat = 1.0
        self.tau_hat = 100.0
        self.theta_hat = self.d * self.dt

        self.t = 0

    def _estimate_delay(self):
        y = np.array(self.y_hist, dtype=float)
        u = np.array(self.u_hist, dtype=float)
        if len(y) < 10 or len(u) < 10:
            return self.d

        n = min(len(y), self.corr_win)
        yw = y[-n:]
        uw = u[-n:]

        dy = np.diff(yw, prepend=yw[0])

        max_lag = min(self.d_max, n-2)
        if max_lag <= 1:
            return self.d

        corrs = []
        for lag in range(0, max_lag+1):
            u_lag = np.roll(uw, shift=lag)
            valid = n - lag
            if valid < 5:
                break
            c = np.corrcoef(dy[lag:], u_lag[lag:])[0,1] if np.std(dy[lag:])>1e-9 and np.std(u_lag[lag:])>1e-9 else 0.0
            corrs.append(c)

        if not corrs:
            return self.d

        best_lag = int(np.nanargmax(corrs))
        if abs(corrs[best_lag]) < 0.05:
            return self.d

        step_limit = 5
        new_d = int(np.clip(best_lag, self.d - step_limit, self.d + step_limit))
        return new_d

    def update(self, u_k, y_k):
        self.t += 1
        self.u_hist.append(float(u_k))
        self.y_hist.append(float(y_k))

        if len(self.y_hist) < 2 or len(self.u_hist) <= self.d:
            return None

        y_k_1 = self.y_hist[-2]
        u_d = self.u_hist[-(self.d+1)]
        phi = np.array([y_k_1, u_d]).reshape(2,1)

        P_phi = self.P @ phi
        denom = self.rls_lambda + (phi.T @ P_phi)[0,0]
        K_gain = P_phi / denom
        y_hat = float((self.theta.reshape(1,2) @ phi)[0,0])
        err = y_k - y_hat
        self.theta = self.theta + (K_gain.flatten() * err)
        self.P = (self.P - K_gain @ phi.T @ self.P) / self.rls_lambda

        if (self.t % self.retune_period) == 0:
            self.d = self._estimate_delay()
            a = float(self.theta[0])
            b = float(self.theta[1])

            a = float(np.clip(a, 1e-4, 0.9999))
            if a <= 0.0 or a >= 1.0:
                return None

            tau_c = -self.dt / np.log(a)
            K_c = b / (1.0 - a + 1e-9)
            theta_c = max(0.0, self.d * self.dt)

            self.K_hat     = (1.0 - self.smooth_alpha) * self.K_hat     + self.smooth_alpha * float(np.clip(K_c, 0.05, 50.0))
            self.tau_hat   = (1.0 - self.smooth_alpha) * self.tau_hat   + self.smooth_alpha * float(np.clip(tau_c, 1.0, 5000.0))
            self.theta_hat = (1.0 - self.smooth_alpha) * self.theta_hat + self.smooth_alpha * float(np.clip(theta_c, 0.0, 1000.0))

            self.estimates.append((self.t, self.K_hat, self.tau_hat, self.theta_hat))
            return self.K_hat, self.tau_hat, self.theta_hat

        return None


# ========== HVAC控制器 ==========
class HVACController:
    def __init__(self, season="winter", dt=1.0, init_params=(0.8, 60.0, 5.0, 40.0)):
        self.season = season
        K0, tau0, theta0, lam0 = init_params
        self.pid = IMCPID(K=K0, tau=tau0, theta=theta0, lam=lam0, dt=dt)

        self.min_on_time  = 300
        self.min_off_time = 300
        self.hysteresis = 0.3
        self.emergency_shutdown_threshold = 2.0

        self.last_switch_time = 0
        self.current_state = "OFF"
        self.time_elapsed = 0

    def control(self, setTemp, roomTemp):
        pid_output = self.pid.compute(setTemp, roomTemp)
        self.time_elapsed += 1

        if roomTemp > setTemp + self.emergency_shutdown_threshold:
            self.current_state = "OFF"
            self.last_switch_time = self.time_elapsed
            return "HEATER OFF (EMERGENCY)", pid_output

        temp_diff = setTemp - roomTemp
        need_heat = temp_diff > self.hysteresis
        stop_heat = temp_diff < -self.hysteresis

        time_since_switch = self.time_elapsed - self.last_switch_time
        if self.current_state == "ON":
            if stop_heat and time_since_switch >= self.min_on_time:
                self.current_state = "OFF"
                self.last_switch_time = self.time_elapsed
        else:
            if need_heat and time_since_switch >= self.min_off_time:
                self.current_state = "ON"
                self.last_switch_time = self.time_elapsed

        return f"HEATER {self.current_state}", pid_output


# ========== 多种房屋热模型 ==========
class HouseModel:
    """基准模型"""
    def __init__(self, initial_temp=18.0):
        self.temp = initial_temp
        self.outdoor_temp = 10.0
        self.heat_loss_coeff = 0.01
        self.heating_power = 0.5
        self.thermal_mass = 100.0

    def update(self, heating_on):
        temp_diff = self.temp - self.outdoor_temp
        heat_loss = self.heat_loss_coeff * temp_diff
        heat_gain = self.heating_power if heating_on else 0.0
        delta_temp = (heat_gain - heat_loss) / self.thermal_mass
        self.temp += delta_temp
        self.temp += np.random.normal(0, 0.005)
        return self.temp


class LightHouseModel(HouseModel):
    """轻型房屋：升温快，散热快"""
    def __init__(self, initial_temp=18.0):
        super().__init__(initial_temp)
        self.heat_loss_coeff = 0.03   # 漏热快
        self.heating_power = 0.6      # 加热稍强
        self.thermal_mass = 40.0      # 热惯性小


class HeavyHouseModel(HouseModel):
    """重型房屋：升温慢，保温好"""
    def __init__(self, initial_temp=18.0):
        super().__init__(initial_temp)
        self.heat_loss_coeff = 0.005  # 漏热慢
        self.heating_power = 0.5
        self.thermal_mass = 300.0     # 热惯性大


class LeakyHouseModel(HouseModel):
    """漏风房屋：极易散热"""
    def __init__(self, initial_temp=18.0):
        super().__init__(initial_temp)
        self.heat_loss_coeff = 0.05   # 极易散热
        self.heating_power = 1.0      # 加热强
        self.thermal_mass = 100.0


# ========== 仿真（选择房屋模型） ==========
np.random.seed(7)
dt = 1.0
controller = HVACController(season="winter", dt=dt, init_params=(0.8, 60.0, 5.0, 40.0))

# 切换不同房屋模型: HouseModel, LightHouseModel, HeavyHouseModel, LeakyHouseModel
house = HouseModel(initial_temp=18.0)

identifier = OnlineRLSFOPDT(
    dt=dt,
    d_max=300,
    lam_corr=2000,
    rls_lambda=0.995,
    init_var=1e3,
    retune_period=120,
    smooth_alpha=0.25
)

setTemp = 22.0
steps = 10000

temps, actions, outputs = [], [], []
K_trace, tau_trace, theta_trace, t_trace = [], [], [], []

for t in range(steps):
    action, pid_output = controller.control(setTemp, house.temp)
    heating_on = 1 if "ON" in action else 0

    est = identifier.update(heating_on, house.temp)
    if est is not None:
        K_hat, tau_hat, theta_hat = est
        lam_new = 40.0
        controller.pid.set_params(K_hat, tau_hat, theta_hat, lam=lam_new)

    roomTemp = house.update(heating_on)

    temps.append(roomTemp)
    actions.append(heating_on)
    outputs.append(pid_output)

    t_trace.append(t)
    K_trace.append(identifier.K_hat)
    tau_trace.append(identifier.tau_hat)
    theta_trace.append(identifier.theta_hat)

# ========== 绘图 ==========
fig = plt.figure(figsize=(14, 12))

ax1 = plt.subplot(4,1,1)
ax1.plot(range(steps), temps, label="Room Temperature (°C)", color="tab:blue")
ax1.axhline(setTemp, color="tab:red", linestyle="--", label="Set Temperature")
ax1.axhline(setTemp + controller.hysteresis, color="tab:orange", linestyle=":", label="Hysteresis Band")
ax1.axhline(setTemp - controller.hysteresis, color="tab:orange", linestyle=":")
ax1.axhline(setTemp + controller.emergency_shutdown_threshold, color="tab:red", linestyle="-.", label="Emergency Shutdown")
ax1.set_ylabel("Temperature (°C)")
ax1.set_title("HVAC Control with Online RLS Identification + Adaptive IMC-PID")
ax1.legend(loc="upper right")
ax1.grid(True)

ax2 = plt.subplot(4,1,2)
ax2.plot(range(steps), outputs, label="PID Output", color="tab:purple")
ax2.axhline(0, color="black", linestyle="--")
ax2.set_ylabel("PID Output")
ax2.legend(loc="upper right")
ax2.grid(True)

ax3 = plt.subplot(4,1,3)
ax3.step(range(steps), actions, label="Heater State (1=ON,0=OFF)",
         color="tab:green", where="post")
ax3.set_xlabel("Time step")
ax3.set_ylabel("Control Signal")
ax3.legend(loc="upper right")
ax3.grid(True)

ax4 = plt.subplot(4,1,4)
ax4.plot(t_trace, K_trace, label="K_hat")
ax4.plot(t_trace, tau_trace, label="tau_hat")
ax4.plot(t_trace, theta_trace, label="theta_hat")
ax4.set_ylabel("Estimated Params")
ax4.set_xlabel("Time step")
ax4.legend(loc="upper right")
ax4.grid(True)

plt.tight_layout()
plt.show()
